from PyQt5.QtCore import QThread, pyqtSignal
from socket import AF_INET, SOCK_STREAM, socket
from Data_Pack import myPack, myRecv
from torch import save as T_save, load as T_load
from torch.utils.tensorboard import SummaryWriter


class FC_HDRM_Work(QThread):
    sig_confin = pyqtSignal()
    sig_dsfin = pyqtSignal()
    sig_samfin = pyqtSignal()
    sig_wffin = pyqtSignal()
    sig_nntfin = pyqtSignal()
    sig_nnsfin = pyqtSignal()
    sig_nnvcfpfin = pyqtSignal()
    sig_nngsfin = pyqtSignal()
    sig_nngdfin = pyqtSignal()
    sig_logfin = pyqtSignal()

    def __init__(self):
        super(FC_HDRM_Work, self).__init__()
        # 状态变量 - 线程控制
        self.thr_state = 0
        self.thr_msg = None
        # 状态变量 - 网络连接
        self.tcp_on = False
        self.tcp_ip = None
        self.tcp_port = None
        self.tcp_text = None
        self.sc_c = socket(AF_INET, SOCK_STREAM)
        self.wsec_c = 120.0
        self.sc_c.settimeout(self.wsec_c)
        # 状态变量 - 数据读取
        self.dset_on = False
        self.cfp_path = ''
        self.rcfp_len = None
        # 状态变量 - 样本显示
        self.sam_ok = False
        self.sam_idx = None
        self.sam_img = None
        self.sam_a_idx = 0
        self.sam_p_idx = 0
        self.sam_n_idx = 0
        self.sam_p_str = ''
        self.sam_n_str = ''
        # 状态变量 - 权重文件
        self.nn_wsave = False
        self.nn_wfok = False
        self.nn_wpath = ''
        self.nn_wfile = None
        # 状态变量 - 网络运行参数
        self.nnt_epoch = 1
        self.nnt_ok = False
        self.nns_ok = False
        self.nnvcfp_ok = False
        # 状态变量 - 网络状态
        self.nn_workon = False
        self.nn_mode = 1
        self.nn_nowepoch = 0
        self.nn_nowbatch = 0
        self.nn_aok = False
        self.nn_lok = False
        # 状态变量 - 网络数据
        self.a_num = None
        self.a_cor = None
        self.a_rate = None
        self.a_losses = None
        self.a_depoch = 1
        # 状态变量 - 数据记录
        self.log_dir = ''
        self.log_sw = None

    def run(self):
        # 0:初始化状态
        # 1:网络连接
        # 2:数据读取
        # 3:样本展示
        # 4:存取权重文件
        # 5:请求神经网络训练
        # 6:请求神经网络中止
        # 7:请求神经网络样本验证
        # 8:请求查看网络状态
        # 9:请求获取网络数据
        # 10:数据记录
        if self.thr_state == 1:
            self.Ser_con()
            self.sig_confin.emit()
        elif self.thr_state == 2:
            self.Dset_init()
            self.sig_dsfin.emit()
        elif self.thr_state == 3:
            self.Sam_get()
            self.sig_samfin.emit()
        elif self.thr_state == 4:
            self.Wfile_sl()
            self.sig_wffin.emit()
        elif self.thr_state == 5:
            self.NN_train()
            self.sig_nntfin.emit()
        elif self.thr_state == 6:
            self.NN_stop()
            self.sig_nnsfin.emit()
        elif self.thr_state == 7:
            self.NN_vcfp()
            self.sig_nnvcfpfin.emit()
        elif self.thr_state == 8:
            self.NN_gs()
            self.sig_nngsfin.emit()
        elif self.thr_state == 9:
            self.NN_gd()
            self.sig_nngdfin.emit()
        elif self.thr_state == 10:
            self.Log_sw()
            self.sig_logfin.emit()

    def Ser_con(self):
        if self.tcp_on:
            self.tcp_text = {'cmd': 'close'}
            self.sc_c.sendall(myPack(self.tcp_text))
            self.sc_c.close()
            self.sc_c = socket(AF_INET, SOCK_STREAM)
            self.sc_c.settimeout(self.wsec_c)
            self.thr_msg = '已断开与服务端的连接'
            self.tcp_on = False
            self.dset_on = False
            self.nn_workon = False
        else:
            try:
                self.sc_c.connect((self.tcp_ip, self.tcp_port))
            except Exception as e:
                self.thr_msg = e
                self.tcp_on = False
            else:
                self.thr_msg = '与服务端已建立连接'
                self.tcp_on = True

    def Dset_init(self):
        self.tcp_text = {'cmd': 'loaddset', 'cfp_path': self.cfp_path}
        self.sc_c.sendall(myPack(self.tcp_text))
        try:
            self.tcp_text = myRecv(self.sc_c)
            self.dset_on = self.tcp_text['dset_on']
        except Exception as e:
            self.thr_msg = e
            self.dset_on = False
        else:
            if self.dset_on:
                self.thr_msg = '数据集载入成功！'
                self.rcfp_len = self.tcp_text['rcfp_len']
            else:
                self.thr_msg = '服务端异常'

    def Sam_get(self):
        self.tcp_text = {'cmd': 'samget',
                        'sam_idx': self.sam_idx}
        self.sc_c.sendall(myPack(self.tcp_text))
        try:
            self.tcp_text = myRecv(self.sc_c)
            self.sam_img = self.tcp_text['sam_img']
            self.sam_a_idx = self.tcp_text['sam_a_idx']
            self.sam_p_idx = self.tcp_text['sam_p_idx']
            self.sam_n_idx = self.tcp_text['sam_n_idx']
            self.sam_p_str = self.tcp_text['sam_p_str']
            self.sam_n_str = self.tcp_text['sam_n_str']
        except Exception as e:
            self.thr_msg = e
            self.sam_ok = False
        else:
            self.sam_img = self.sam_img.toqpixmap()
            self.sam_ok = True

    def Wfile_sl(self):
        self.tcp_text = {'cmd': 'slwfile', 'nn_wsave': self.nn_wsave}
        try:
            if not self.nn_wsave:
                self.nn_wfile = T_load(self.nn_wpath)
                self.tcp_text['nn_wfile'] = self.nn_wfile
            self.sc_c.sendall(myPack(self.tcp_text))
            self.tcp_text = myRecv(self.sc_c)
            if self.nn_wsave:
                self.nn_wfile = self.tcp_text['nn_wfile']
                T_save(self.nn_wfile, self.nn_wpath)
            self.nn_wfok = self.tcp_text['nn_wfok']
        except Exception as e:
            self.thr_msg = e
            self.nn_wfok = False
        else:
            if self.nn_wsave:
                self.thr_msg = '权重文件保存成功！'
            else:
                self.thr_msg = '权重文件读取成功！'

    def NN_train(self):
        self.tcp_text = {'cmd': 'nnt',
                        'nnt_epoch': self.nnt_epoch}
        self.sc_c.sendall(myPack(self.tcp_text))
        try:
            self.tcp_text = myRecv(self.sc_c)
            self.nnt_ok = self.tcp_text['nnt_ok']
            self.nn_workon = self.tcp_text['nn_workon']
            self.nn_mode = self.tcp_text['nn_mode']
        except Exception as e:
            self.thr_msg = e
            self.nnt_ok = False
        else:
            if self.nnt_ok:
                self.thr_msg = '神经网络已开始训练！'
            else:
                self.thr_msg = '神经网络已在工作中'

    def NN_stop(self):
        self.tcp_text = {'cmd': 'nns'}
        self.sc_c.sendall(myPack(self.tcp_text))
        try:
            self.tcp_text = myRecv(self.sc_c)
            self.nns_ok = self.tcp_text['nns_ok']
            self.nn_workon = self.tcp_text['nn_workon']
            self.nn_mode = self.tcp_text['nn_mode']
        except Exception as e:
            self.thr_msg = e
            self.nns_ok = False
        else:
            if self.nns_ok:
                self.thr_msg = '神经网络已中止！'
            else:
                self.thr_msg = '神经网络未在工作中'

    def NN_vcfp(self):
        self.tcp_text = {'cmd': 'nnvcfp'}
        self.sc_c.sendall(myPack(self.tcp_text))
        try:
            self.tcp_text = myRecv(self.sc_c)
            self.nnvcfp_ok = self.tcp_text['nnvcfp_ok']
            self.nn_workon = self.tcp_text['nn_workon']
            self.nn_mode = self.tcp_text['nn_mode']
        except Exception as e:
            self.thr_msg = e
            self.nnvcfp_ok = False
        else:
            if self.nnvcfp_ok:
                self.thr_msg = '神经网络已开始CFP验证！'
            else:
                self.thr_msg = '神经网络已在工作中'

    def NN_gs(self):
        self.tcp_text = {'cmd': 'nngs'}
        self.sc_c.sendall(myPack(self.tcp_text))
        try:
            self.tcp_text = myRecv(self.sc_c)
            self.nn_workon = self.tcp_text['nn_workon']
            self.nn_mode = self.tcp_text['nn_mode']
            self.nn_nowepoch = self.tcp_text['nn_nowepoch']
            self.nn_nowbatch = self.tcp_text['nn_nowbatch']
        except Exception as e:
            self.thr_msg = e
        else:
            self.thr_msg = '工作中: {}, 模式: {},epoch: {:0>4d}, batch: {:0>4d}'.format(
                self.nn_workon, self.nn_mode, self.nn_nowepoch + 1, self.nn_nowbatch + 1)

    def NN_gd(self):
        self.tcp_text = {'cmd': 'nngd'}
        self.sc_c.sendall(myPack(self.tcp_text))
        try:
            self.tcp_text = myRecv(self.sc_c)
            self.nn_aok = self.tcp_text['nn_aok']
            if self.nn_aok:
                self.a_num = self.tcp_text['a_num']
                self.a_cor = self.tcp_text['a_cor']
                self.a_rate = self.tcp_text['a_rate']
            self.nn_lok = self.tcp_text['nn_lok']
            if self.nn_lok:
                self.a_losses = self.tcp_text['a_losses']
            self.a_depoch = self.tcp_text['a_depoch']
        except Exception as e:
            self.thr_msg = e
        else:
            self.thr_msg = 'Accuracy数据: {}, Loss数据: {}'.format(
                                        self.nn_aok, self.nn_lok)

    def Log_sw(self):
        self.log_sw = SummaryWriter(self.log_dir)
        if self.nn_lok:
            for i in range(self.a_depoch):
                self.log_sw.add_scalar('Triplet Loss', self.a_losses[i], i + 1)
        self.log_sw.flush()
        self.log_sw.close()
